# Copyright 2025 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Generate data for test."""
import numpy as np

def get_init_params(vocab_size, hidden_size):
    # Generate initialization parameters
    np.random.seed(42)
    weight_shape = (vocab_size, hidden_size)
    return {
        "input_ids": 10 * np.random.rand(2, 2),
        "position_ids": 10 * np.random.rand(2, 2),
        "tokentype_ids": 10 * np.random.rand(2, 2),
        "weight": 0.01 * np.random.randn(*weight_shape),
    }

def get_golden() -> dict[str, np.ndarray]:
    """Generate golden data for test."""
    output_only = np.array(
        [[[5.46264648e-03, -2.02941895e-03, -2.18200684e-03, 1.09863281e-02,
           8.23974609e-03, 8.11767578e-03, 1.30615234e-02, 2.09808350e-04,
           6.80541992e-03, -3.09753418e-03, 3.23486328e-03, -1.30462646e-03,
           9.68933105e-04, 5.95092773e-03, -8.17871094e-03, 2.08740234e-02,
           -1.00708008e-02, -1.21459961e-02, 1.15966797e-02, 7.93457031e-03,
           6.25610352e-03, 6.28662109e-03, -1.22070312e-04, -8.97216797e-03,
           7.59124756e-04, -6.77490234e-03, 9.76562500e-03, -1.47247314e-03,
           -8.23974609e-03, -3.21960449e-03, 4.11987305e-03, -5.64575195e-03,
           -8.23974609e-03, 2.44140625e-03, 2.45666504e-03, -5.06591797e-03,
           -4.69970703e-03, 2.31933594e-03, -1.44653320e-02, -1.40991211e-02,
           -7.17163086e-03, -2.13623047e-03, 3.11279297e-03, 1.47705078e-02,
           8.60595703e-03, -1.60217285e-03, -1.89781189e-04, -1.00097656e-02,
           -1.85012817e-04, -2.88391113e-03, 3.21960449e-03, -8.30078125e-03,
           5.18798828e-03, 1.53198242e-02, -1.09100342e-03, 4.02832031e-03,
           6.89697266e-03, -3.99780273e-03, 2.24304199e-03, 1.25885010e-04,
           9.76562500e-04, -7.72094727e-03, 2.46047974e-04, 4.97436523e-03,
           1.45263672e-02, 9.58251953e-03, 2.14843750e-02, -7.65991211e-03,
           8.72802734e-03, 1.83105469e-03, 2.18505859e-02, -8.05664062e-03,
           -8.42285156e-03, -5.98144531e-03, -2.12402344e-02, -5.24902344e-03,
           -7.59887695e-03, 1.50299072e-03, 3.41796875e-03, 1.87988281e-02,
           9.52148438e-03, -5.76782227e-03, -8.97216797e-03, 4.91333008e-03,
           -1.31835938e-02, 1.83105469e-02, 1.17797852e-02, -4.69970703e-03,
           -1.70898438e-02, 1.35498047e-02, -1.14440918e-03, 1.23901367e-02,
           -1.59912109e-02, -5.98144531e-03, 5.24520874e-05, 4.69207764e-04,
           -4.48608398e-03, 6.22558594e-03, -1.06811523e-02, -1.42669678e-03],
          [-2.08740234e-02, 1.69677734e-02, 2.10571289e-03, -9.68933105e-04,
           -5.46264648e-03, 3.99780273e-03, -3.75747681e-04, 1.10473633e-02,
           1.14440918e-03, 1.50299072e-03, -3.63159180e-03, -5.68389893e-04,
           3.08227539e-03, -1.70898438e-02, -1.34887695e-02, 7.44628906e-03,
           1.70898438e-03, -1.83868408e-03, 1.84059143e-04, 3.47900391e-03,
           -5.40161133e-03, -7.78198242e-03, 1.95312500e-03, -9.76562500e-03,
           4.08935547e-03, -1.69677734e-02, 1.03149414e-02, 4.73022461e-03,
           2.56347656e-03, 9.82666016e-03, 1.66015625e-02, 1.01318359e-02,
           -1.84326172e-02, -1.28173828e-02, -6.25610352e-03, 2.61306763e-04,
           5.18798828e-03, -7.26318359e-03, 1.86920166e-03, -7.56835938e-03,
           -6.10351562e-03, -1.40380859e-02, -9.21630859e-03, -1.34887695e-02,
           -9.76562500e-03, 1.05590820e-02, -9.52148438e-03, 2.63671875e-02,
           4.94384766e-03, 1.84631348e-03, -8.60595703e-03, 6.98852539e-03,
           -5.76782227e-03, 1.22070312e-03, 2.56347656e-02, -9.61303711e-04,
           1.14746094e-02, -7.01904297e-03, -3.49044800e-04, 1.77001953e-02,
           -6.25610352e-03, 1.80664062e-02, 7.08007812e-03, -5.61523438e-03,
           6.31713867e-03, 9.70458984e-03, 6.22558594e-03, -1.57470703e-02,
           -7.26318359e-03, -2.47192383e-03, -7.43865967e-04, 6.19506836e-03,
           1.77764893e-03, -1.33666992e-02, 3.79943848e-03, 6.10351562e-03,
           5.58471680e-03, 1.08032227e-02, 8.36181641e-03, 4.57763672e-03,
           -7.01904297e-04, -1.66015625e-02, 4.30297852e-03, 2.07519531e-03,
           2.71606445e-03, -1.27563477e-02, -1.08032227e-02, 1.05590820e-02,
           -3.94821167e-04, 6.80541992e-03, 2.82287598e-04, 2.97546387e-04,
           9.39941406e-03, -5.15747070e-03, 9.61303711e-04, -4.60815430e-03,
           -4.33349609e-03, -3.09753418e-03, 2.22778320e-03, -4.79125977e-03]],
         [[6.46972656e-03, 2.16064453e-02, -3.08227539e-03, 2.19726562e-03,
           2.48718262e-03, 1.57470703e-02, -9.53674316e-04, 2.79235840e-03,
           6.07299805e-03, 1.86920166e-03, -4.45556641e-03, 1.93786621e-03,
           1.07421875e-02, -1.02539062e-02, 1.32751465e-03, -6.98852539e-03,
           1.19628906e-02, -1.52587891e-02, -5.58471680e-03, 3.76892090e-03,
           1.56250000e-02, -6.56127930e-04, -5.55419922e-03, 1.87988281e-02,
           -1.44653320e-02, -2.19726562e-02, 4.39453125e-03, -5.03540039e-03,
           -1.01928711e-02, 7.08007812e-03, 2.44140625e-03, -5.64575195e-03,
           -1.28173828e-02, 8.72802734e-03, 6.50024414e-03, -9.91821289e-04,
           1.84326172e-02, -1.06811523e-02, -1.52587891e-02, -6.92749023e-03,
           -4.55856323e-04, 2.42614746e-03, -2.41088867e-03, 3.52478027e-03,
           -1.25122070e-02, 1.44653320e-02, -8.20159912e-04, 1.11694336e-02,
           3.43322754e-03, 4.57763672e-03, 5.70678711e-03, 4.48608398e-03,
           6.43920898e-03, 1.33056641e-02, 1.96838379e-03, 7.08007812e-03,
           -8.96453857e-04, 1.44042969e-02, -6.77490234e-03, 1.80664062e-02,
           -4.02450562e-04, -1.42822266e-02, 1.28173828e-03, -6.80541992e-03,
           8.42285156e-03, -6.53076172e-03, -4.45556641e-03, -1.89208984e-02,
           -4.51660156e-03, -2.42919922e-02, -1.58691406e-02, 7.59887695e-03,
           7.87353516e-03, 4.24194336e-03, -9.64355469e-03, -4.76837158e-04,
           -3.60012054e-05, -1.15966797e-02, 1.50146484e-02, 8.78906250e-03,
           -2.21252441e-03, 2.68936157e-04, 2.09045410e-03, -2.03857422e-02,
           -2.47192383e-03, -6.80541992e-03, -1.00097656e-02, -2.80761719e-03,
           1.79443359e-02, 6.40869141e-03, -5.70678711e-03, 5.73730469e-03,
           1.39770508e-02, 9.21630859e-03, 5.95092773e-04, -6.46972656e-03,
           6.98852539e-03, 3.93676758e-03, 8.97216797e-03, 6.34765625e-03],
          [4.84466553e-04, -8.30078125e-03, 2.70080566e-03, -5.03540039e-04,
           -2.39562988e-03, -9.09423828e-03, -5.76782227e-03, 7.56835938e-03,
           5.00488281e-03, -9.76562500e-03, 9.91821289e-04, 7.50732422e-03,
           -1.67236328e-02, 5.43212891e-03, -6.62231445e-03, 5.70678711e-03,
           -7.62939453e-03, -1.80664062e-02, -1.62353516e-02, 4.80651855e-04,
           2.59399414e-03, -9.03320312e-03, 6.37817383e-03, -1.66015625e-02,
           -6.59942627e-04, -1.20849609e-02, -6.53076172e-03, 4.74929810e-04,
           -8.60595703e-03, -3.84521484e-03, 1.00708008e-02, -5.76782227e-03,
           8.36181641e-03, -1.12915039e-02, 5.31005859e-03, 1.44042969e-02,
           -2.46582031e-02, -7.99560547e-03, 5.76782227e-03, -2.02941895e-03,
           3.70788574e-03, -6.04248047e-03, 8.65936279e-04, -1.55639648e-03,
           1.16577148e-02, 2.54821777e-03, 3.37219238e-03, -4.11987305e-03,
           -4.88281250e-03, -4.33349609e-03, 3.93676758e-03, -4.21142578e-03,
           2.89916992e-03, 2.07519531e-02, 8.72802734e-03, -3.26538086e-03,
           1.20239258e-02, -4.08935547e-03, -2.03857422e-02, -1.00708008e-02,
           -1.86767578e-02, -3.50952148e-03, 1.84059143e-04, 1.67236328e-02,
           3.26538086e-03, -2.19726562e-03, 8.30078125e-03, -2.20947266e-02,
           2.34985352e-03, 7.72094727e-03, -1.47705078e-02, 1.14135742e-02,
           3.38745117e-03, -4.15039062e-03, 6.31713867e-03, 2.27050781e-02,
           1.81579590e-03, 2.48718262e-03, -4.60815430e-03, -8.48388672e-03,
           8.30078125e-03, -8.54492188e-03, 7.17163086e-04, -4.79125977e-03,
           4.79125977e-03, 3.34167480e-03, 1.03759766e-02, -5.09643555e-03,
           -2.70080566e-03, -9.76562500e-03, -4.45556641e-03, 3.76892090e-03,
           7.56835938e-03, -9.21630859e-03, 8.66699219e-03, 1.35498047e-02,
           4.11987305e-03, 1.87988281e-02, -7.75146484e-03, -1.24511719e-02]]])
    return {
        "output_only": output_only,
    }

def get_gpu_datas() -> dict[str, np.ndarray]:
    """Generate gpu data for test."""
    output_only = np.array(
        [[[5.46264648e-03, -2.02941895e-03, -2.18200684e-03, 1.09863281e-02,
           8.23974609e-03, 8.11767578e-03, 1.30615234e-02, 2.09808350e-04,
           6.80541992e-03, -3.09753418e-03, 3.23486328e-03, -1.30462646e-03,
           9.68933105e-04, 5.95092773e-03, -8.17871094e-03, 2.08740234e-02,
           -1.00708008e-02, -1.21459961e-02, 1.15966797e-02, 7.93457031e-03,
           6.25610352e-03, 6.28662109e-03, -1.22070312e-04, -8.97216797e-03,
           7.59124756e-04, -6.77490234e-03, 9.76562500e-03, -1.47247314e-03,
           -8.23974609e-03, -3.21960449e-03, 4.11987305e-03, -5.64575195e-03,
           -8.23974609e-03, 2.44140625e-03, 2.45666504e-03, -5.06591797e-03,
           -4.69970703e-03, 2.31933594e-03, -1.44653320e-02, -1.40991211e-02,
           -7.17163086e-03, -2.13623047e-03, 3.11279297e-03, 1.47705078e-02,
           8.60595703e-03, -1.60217285e-03, -1.89781189e-04, -1.00097656e-02,
           -1.85012817e-04, -2.88391113e-03, 3.21960449e-03, -8.30078125e-03,
           5.18798828e-03, 1.53198242e-02, -1.09100342e-03, 4.02832031e-03,
           6.89697266e-03, -3.99780273e-03, 2.24304199e-03, 1.25885010e-04,
           9.76562500e-04, -7.72094727e-03, 2.46047974e-04, 4.97436523e-03,
           1.45263672e-02, 9.58251953e-03, 2.14843750e-02, -7.65991211e-03,
           8.72802734e-03, 1.83105469e-03, 2.18505859e-02, -8.05664062e-03,
           -8.42285156e-03, -5.98144531e-03, -2.12402344e-02, -5.24902344e-03,
           -7.59887695e-03, 1.50299072e-03, 3.41796875e-03, 1.87988281e-02,
           9.52148438e-03, -5.76782227e-03, -8.97216797e-03, 4.91333008e-03,
           -1.31835938e-02, 1.83105469e-02, 1.17797852e-02, -4.69970703e-03,
           -1.70898438e-02, 1.35498047e-02, -1.14440918e-03, 1.23901367e-02,
           -1.59912109e-02, -5.98144531e-03, 5.24520874e-05, 4.69207764e-04,
           -4.48608398e-03, 6.22558594e-03, -1.06811523e-02, -1.42669678e-03],
          [-2.08740234e-02, 1.69677734e-02, 2.10571289e-03, -9.68933105e-04,
           -5.46264648e-03, 3.99780273e-03, -3.75747681e-04, 1.10473633e-02,
           1.14440918e-03, 1.50299072e-03, -3.63159180e-03, -5.68389893e-04,
           3.08227539e-03, -1.70898438e-02, -1.34887695e-02, 7.44628906e-03,
           1.70898438e-03, -1.83868408e-03, 1.84059143e-04, 3.47900391e-03,
           -5.40161133e-03, -7.78198242e-03, 1.95312500e-03, -9.76562500e-03,
           4.08935547e-03, -1.69677734e-02, 1.03149414e-02, 4.73022461e-03,
           2.56347656e-03, 9.82666016e-03, 1.66015625e-02, 1.01318359e-02,
           -1.84326172e-02, -1.28173828e-02, -6.25610352e-03, 2.61306763e-04,
           5.18798828e-03, -7.26318359e-03, 1.86920166e-03, -7.56835938e-03,
           -6.10351562e-03, -1.40380859e-02, -9.21630859e-03, -1.34887695e-02,
           -9.76562500e-03, 1.05590820e-02, -9.52148438e-03, 2.63671875e-02,
           4.94384766e-03, 1.84631348e-03, -8.60595703e-03, 6.98852539e-03,
           -5.76782227e-03, 1.22070312e-03, 2.56347656e-02, -9.61303711e-04,
           1.14746094e-02, -7.01904297e-03, -3.49044800e-04, 1.77001953e-02,
           -6.25610352e-03, 1.80664062e-02, 7.08007812e-03, -5.61523438e-03,
           6.31713867e-03, 9.70458984e-03, 6.22558594e-03, -1.57470703e-02,
           -7.26318359e-03, -2.47192383e-03, -7.43865967e-04, 6.19506836e-03,
           1.77764893e-03, -1.33666992e-02, 3.79943848e-03, 6.10351562e-03,
           5.58471680e-03, 1.08032227e-02, 8.36181641e-03, 4.57763672e-03,
           -7.01904297e-04, -1.66015625e-02, 4.30297852e-03, 2.07519531e-03,
           2.71606445e-03, -1.27563477e-02, -1.08032227e-02, 1.05590820e-02,
           -3.94821167e-04, 6.80541992e-03, 2.82287598e-04, 2.97546387e-04,
           9.39941406e-03, -5.15747070e-03, 9.61303711e-04, -4.60815430e-03,
           -4.33349609e-03, -3.09753418e-03, 2.22778320e-03, -4.79125977e-03]],
         [[6.46972656e-03, 2.16064453e-02, -3.08227539e-03, 2.19726562e-03,
           2.48718262e-03, 1.57470703e-02, -9.53674316e-04, 2.79235840e-03,
           6.07299805e-03, 1.86920166e-03, -4.45556641e-03, 1.93786621e-03,
           1.07421875e-02, -1.02539062e-02, 1.32751465e-03, -6.98852539e-03,
           1.19628906e-02, -1.52587891e-02, -5.58471680e-03, 3.76892090e-03,
           1.56250000e-02, -6.56127930e-04, -5.55419922e-03, 1.87988281e-02,
           -1.44653320e-02, -2.19726562e-02, 4.39453125e-03, -5.03540039e-03,
           -1.01928711e-02, 7.08007812e-03, 2.44140625e-03, -5.64575195e-03,
           -1.28173828e-02, 8.72802734e-03, 6.50024414e-03, -9.91821289e-04,
           1.84326172e-02, -1.06811523e-02, -1.52587891e-02, -6.92749023e-03,
           -4.55856323e-04, 2.42614746e-03, -2.41088867e-03, 3.52478027e-03,
           -1.25122070e-02, 1.44653320e-02, -8.20159912e-04, 1.11694336e-02,
           3.43322754e-03, 4.57763672e-03, 5.70678711e-03, 4.48608398e-03,
           6.43920898e-03, 1.33056641e-02, 1.96838379e-03, 7.08007812e-03,
           -8.96453857e-04, 1.44042969e-02, -6.77490234e-03, 1.80664062e-02,
           -4.02450562e-04, -1.42822266e-02, 1.28173828e-03, -6.80541992e-03,
           8.42285156e-03, -6.53076172e-03, -4.45556641e-03, -1.89208984e-02,
           -4.51660156e-03, -2.42919922e-02, -1.58691406e-02, 7.59887695e-03,
           7.87353516e-03, 4.24194336e-03, -9.64355469e-03, -4.76837158e-04,
           -3.60012054e-05, -1.15966797e-02, 1.50146484e-02, 8.78906250e-03,
           -2.21252441e-03, 2.68936157e-04, 2.09045410e-03, -2.03857422e-02,
           -2.47192383e-03, -6.80541992e-03, -1.00097656e-02, -2.80761719e-03,
           1.79443359e-02, 6.40869141e-03, -5.70678711e-03, 5.73730469e-03,
           1.39770508e-02, 9.21630859e-03, 5.95092773e-04, -6.46972656e-03,
           6.98852539e-03, 3.93676758e-03, 8.97216797e-03, 6.34765625e-03],
          [4.84466553e-04, -8.30078125e-03, 2.70080566e-03, -5.03540039e-04,
           -2.39562988e-03, -9.09423828e-03, -5.76782227e-03, 7.56835938e-03,
           5.00488281e-03, -9.76562500e-03, 9.91821289e-04, 7.50732422e-03,
           -1.67236328e-02, 5.43212891e-03, -6.62231445e-03, 5.70678711e-03,
           -7.62939453e-03, -1.80664062e-02, -1.62353516e-02, 4.80651855e-04,
           2.59399414e-03, -9.03320312e-03, 6.37817383e-03, -1.66015625e-02,
           -6.59942627e-04, -1.20849609e-02, -6.53076172e-03, 4.74929810e-04,
           -8.60595703e-03, -3.84521484e-03, 1.00708008e-02, -5.76782227e-03,
           8.36181641e-03, -1.12915039e-02, 5.31005859e-03, 1.44042969e-02,
           -2.46582031e-02, -7.99560547e-03, 5.76782227e-03, -2.02941895e-03,
           3.70788574e-03, -6.04248047e-03, 8.65936279e-04, -1.55639648e-03,
           1.16577148e-02, 2.54821777e-03, 3.37219238e-03, -4.11987305e-03,
           -4.88281250e-03, -4.33349609e-03, 3.93676758e-03, -4.21142578e-03,
           2.89916992e-03, 2.07519531e-02, 8.72802734e-03, -3.26538086e-03,
           1.20239258e-02, -4.08935547e-03, -2.03857422e-02, -1.00708008e-02,
           -1.86767578e-02, -3.50952148e-03, 1.84059143e-04, 1.67236328e-02,
           3.26538086e-03, -2.19726562e-03, 8.30078125e-03, -2.20947266e-02,
           2.34985352e-03, 7.72094727e-03, -1.47705078e-02, 1.14135742e-02,
           3.38745117e-03, -4.15039062e-03, 6.31713867e-03, 2.27050781e-02,
           1.81579590e-03, 2.48718262e-03, -4.60815430e-03, -8.48388672e-03,
           8.30078125e-03, -8.54492188e-03, 7.17163086e-04, -4.79125977e-03,
           4.79125977e-03, 3.34167480e-03, 1.03759766e-02, -5.09643555e-03,
           -2.70080566e-03, -9.76562500e-03, -4.45556641e-03, 3.76892090e-03,
           7.56835938e-03, -9.21630859e-03, 8.66699219e-03, 1.35498047e-02,
           4.11987305e-03, 1.87988281e-02, -7.75146484e-03, -1.24511719e-02]]])
    return {
        "output_only": output_only,
    }

GOLDEN_DATA = get_golden()
GPU_DATA = get_gpu_datas()
